/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2023-2024. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/*!
 * \file flash_attention_score_grad_tiling_common.cpp
 * \brief
 */

#include "flash_attention_score_grad_tiling_common.h"
#include "log/ops_log.h"

namespace optiling {

ge::graphStatus CheckSoftmaxMaxShape(gert::TilingContext* context, uint32_t b, uint32_t n1, uint32_t s1) {
  auto softmaxMaxShape = context->GetOptionalInputShape(SOFTMAX_MAX);
  if (softmaxMaxShape == nullptr) {
    return ge::GRAPH_SUCCESS;
  }
  auto softmaxMaxShapeDim = softmaxMaxShape->GetStorageShape().GetDimNum();
  if (softmaxMaxShapeDim != 4) {  // softmaxMax only support 4 dimensions
    OPS_LOG_E(context, "The shape of softmaxMax is invalid, got %ld dimensions", softmaxMaxShapeDim);
    return ge::GRAPH_FAILED;
  }
  auto dim0 = softmaxMaxShape->GetStorageShape().GetDim(0);  // 0:b
  auto dim1 = softmaxMaxShape->GetStorageShape().GetDim(1);  // 1:n1
  auto dim2 = softmaxMaxShape->GetStorageShape().GetDim(2);  // 2:s1
  auto dim3 = softmaxMaxShape->GetStorageShape().GetDim(3);  // 3:8
  if (dim0 != b || dim1 != n1 || dim2 != s1 || dim3 != 8) {  // softmaxMax pad to 8
    OPS_LOG_E(context, "The shape of softmaxMax is invalid, got (%ld,%ld,%ld,%ld), should be (%u,%u,%u,8)", dim0, dim1,
              dim2, dim3, b, n1, s1);
    return ge::GRAPH_FAILED;
  }
  return ge::GRAPH_SUCCESS;
}

ge::graphStatus CheckTndSoftmaxMaxShape(gert::TilingContext* context, uint32_t t1, uint32_t n1) {
  auto softmaxMaxShape = context->GetOptionalInputShape(SOFTMAX_MAX);
  if (softmaxMaxShape == nullptr) {
    return ge::GRAPH_SUCCESS;
  }
  auto softmaxMaxShapeDim = softmaxMaxShape->GetStorageShape().GetDimNum();
  if (softmaxMaxShapeDim != 3) {  // TND softmaxMax only support 3 dimensions
    OPS_LOG_E(context, "The shape of softmaxMax is invalid, got %ld dimensions", softmaxMaxShapeDim);
    return ge::GRAPH_FAILED;
  }
  auto dim0 = softmaxMaxShape->GetStorageShape().GetDim(0);  // 0:t1
  auto dim1 = softmaxMaxShape->GetStorageShape().GetDim(1);  // 1:n1
  auto dim2 = softmaxMaxShape->GetStorageShape().GetDim(2);  // 2:8
  if (dim0 != t1 || dim1 != n1 || dim2 != 8) {               // softmaxMax pad to 8
    OPS_LOG_E(context, "The shape of softmaxMax is invalid, got (%ld,%ld,%ld), should be (%u,%u,8)", dim0, dim1, dim2,
              t1, n1);
    return ge::GRAPH_FAILED;
  }
  return ge::GRAPH_SUCCESS;
}

ge::graphStatus CheckSoftmaxSumShape(gert::TilingContext* context, uint32_t b, uint32_t n1, uint32_t s1) {
  auto softmaxSumShape = context->GetOptionalInputShape(SOFTMAX_SUM);
  if (softmaxSumShape == nullptr) {
    return ge::GRAPH_SUCCESS;
  }
  auto softmaxSumShapeDim = softmaxSumShape->GetStorageShape().GetDimNum();
  if (softmaxSumShapeDim != 4) {  // softmaxSum only support 4 dimensions
    OPS_LOG_E(context, "The shape of softmaxSum is invalid, got %ld dimensions", softmaxSumShapeDim);
    return ge::GRAPH_FAILED;
  }
  auto dim0 = softmaxSumShape->GetStorageShape().GetDim(0);  // 0:b
  auto dim1 = softmaxSumShape->GetStorageShape().GetDim(1);  // 1:n1
  auto dim2 = softmaxSumShape->GetStorageShape().GetDim(2);  // 2:s1
  auto dim3 = softmaxSumShape->GetStorageShape().GetDim(3);  // 3:8
  if (dim0 != b || dim1 != n1 || dim2 != s1 || dim3 != 8) {  // softmaxSum pad to 8
    OPS_LOG_E(context, "The shape of softmaxSum is invalid, got (%ld,%ld,%ld,%ld), should be (%u,%u,%u,8)", dim0, dim1,
              dim2, dim3, b, n1, s1);
    return ge::GRAPH_FAILED;
  }
  return ge::GRAPH_SUCCESS;
}

ge::graphStatus CheckTndSoftmaxSumShape(gert::TilingContext* context, uint32_t t1, uint32_t n1) {
  auto softmaxSumShape = context->GetOptionalInputShape(SOFTMAX_SUM);
  if (softmaxSumShape == nullptr) {
    return ge::GRAPH_SUCCESS;
  }
  auto softmaxSumShapeDim = softmaxSumShape->GetStorageShape().GetDimNum();
  if (softmaxSumShapeDim != 3) {  // TND softmaxSum only support 3 dimensions
    OPS_LOG_E(context, "The shape of softmaxSum is invalid, got %ld dimensions", softmaxSumShapeDim);
    return ge::GRAPH_FAILED;
  }
  auto dim0 = softmaxSumShape->GetStorageShape().GetDim(0);  // 0:t1
  auto dim1 = softmaxSumShape->GetStorageShape().GetDim(1);  // 1:n1
  auto dim2 = softmaxSumShape->GetStorageShape().GetDim(2);  // 2:8
  if (dim0 != t1 || dim1 != n1 || dim2 != 8) {               // softmaxSum pad to 8
    OPS_LOG_E(context, "The shape of softmaxSum is invalid, got (%ld,%ld,%ld), should be (%u,%u,8)", dim0, dim1, dim2,
              t1, n1);
    return ge::GRAPH_FAILED;
  }
  return ge::GRAPH_SUCCESS;
}

ge::graphStatus CheckAttentionInShape(gert::TilingContext* context) {
  auto attentionInShape = context->GetOptionalInputShape(ATTENTION_IN);
  if (attentionInShape == nullptr) {
    return ge::GRAPH_SUCCESS;
  }
  auto queryShape = context->GetInputShape(QUERY);
  auto attentionInShapeDim = attentionInShape->GetStorageShape().GetDimNum();
  auto queryShapeDim = queryShape->GetStorageShape().GetDimNum();
  if (attentionInShapeDim != queryShapeDim) {
    OPS_LOG_E(context, "The dimnum of attentionIn %ld should be equal to query %ld", attentionInShapeDim,
              queryShapeDim);
    return ge::GRAPH_FAILED;
  }
  for (size_t i = 0; i < queryShapeDim; i++) {
    if (attentionInShape->GetStorageShape().GetDim(i) != queryShape->GetStorageShape().GetDim(i)) {
      OPS_LOG_E(context, "The dim %zu of attentionIn shape is invalid, got %ld, should be %ld", i,
                attentionInShape->GetStorageShape().GetDim(i), queryShape->GetStorageShape().GetDim(i));
      return ge::GRAPH_FAILED;
    }
  }
  return ge::GRAPH_SUCCESS;
}

ge::graphStatus CheckShapeValid(gert::TilingContext* context, uint32_t b, uint32_t n1, uint32_t s1, uint32_t d) {
  if (context == nullptr) {
    OPS_LOG_E("CheckInnerShape", "context is nullptr");
    return ge::GRAPH_FAILED;
  }

  auto isShapeInValid = (b == 0 || n1 == 0 || s1 == 0 || d == 0);
  if (isShapeInValid) {
    OPS_LOG_E(context, "input shape error, got 0 in bnsd(%u,%u,%u,%u)", b, n1, s1, d);
    return ge::GRAPH_FAILED;
  }

  auto ret = CheckSoftmaxMaxShape(context, b, n1, s1);
  if (ret != ge::GRAPH_SUCCESS) {
    return ret;
  }
  ret = CheckSoftmaxSumShape(context, b, n1, s1);
  if (ret != ge::GRAPH_SUCCESS) {
    return ret;
  }
  ret = CheckAttentionInShape(context);
  if (ret != ge::GRAPH_SUCCESS) {
    return ret;
  }

  return ge::GRAPH_SUCCESS;
}

ge::graphStatus CheckTndShapeValid(gert::TilingContext* context, uint32_t t1, uint32_t n1, uint32_t d) {
  if (context == nullptr) {
    OPS_LOG_E("CheckTndShapeValid", "context is nullptr");
    return ge::GRAPH_FAILED;
  }

  auto isShapeInValid = (t1 == 0 || n1 == 0 || d == 0);
  if (isShapeInValid) {
    OPS_LOG_E(context, "input shape error, got 0 in tnd(%u,%u,%u)", t1, n1, d);
    return ge::GRAPH_FAILED;
  }

  auto ret = CheckTndSoftmaxMaxShape(context, t1, n1);
  if (ret != ge::GRAPH_SUCCESS) {
    return ret;
  }
  ret = CheckTndSoftmaxSumShape(context, t1, n1);
  if (ret != ge::GRAPH_SUCCESS) {
    return ret;
  }
  ret = CheckAttentionInShape(context);
  if (ret != ge::GRAPH_SUCCESS) {
    return ret;
  }

  return ge::GRAPH_SUCCESS;
}

}  // namespace optiling
